{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Finetune and deploy a custom Generative Command model\n",
    "\n",
    "This sample notebook shows you how to finetune and deploy a custom Command model using Amazon SageMaker.\n",
    "\n",
    "> **Note**: This is a reference notebook and it cannot run unless you make changes suggested in the notebook.\n",
    "\n",
    "## Pre-requisites:\n",
    "1. **Note**: This notebook contains elements which render correctly in Jupyter interface. Open this notebook from an Amazon SageMaker Notebook Instance or Amazon SageMaker Studio.\n",
    "1. Ensure that IAM role used has **AmazonSageMakerFullAccess**\n",
    "1. To deploy this ML model successfully, ensure that:\n",
    "    1. Either your IAM role has these three permissions and you have authority to make AWS Marketplace subscriptions in the AWS account used: \n",
    "        1. **aws-marketplace:ViewSubscriptions**\n",
    "        1. **aws-marketplace:Unsubscribe**\n",
    "        1. **aws-marketplace:Subscribe**  \n",
    "    2. or your AWS account has a subscription to the packages for [Cohere Command Finetuning](https://aws.amazon.com/marketplace/pp/prodview-3hsemkam64xmi). If so, skip step: [Subscribe to the finetune algorithm](#1.-Subscribe-to-the-finetune-algorithm)\n",
    "\n",
    "## Contents:\n",
    "1. [Subscribe to the finetune algorithm](#1.-Subscribe-to-the-finetune-algorithm)\n",
    "2. [Finetune Generation Models](#2.-Finetune-the-model)\n",
    "   1. [Upload training and evaluation datasets](#A.-Upload-training-and-evaluation-datasets)\n",
    "   2. [Finetune models on uploaded data](#B.-Finetune-model-on-uploaded-data)\n",
    "3. [Create an endpoint for inference with the custom model](#3.-Create-an-endpoint-for-inference-with-the-custom-model)\n",
    "   1. [Create an endpoint]()\n",
    "   2. [Perform real-time inference]()\n",
    "4. [Clean-up](#4.-Clean-up)\n",
    "    1. [Delete the endpoint](#A.-Delete-the-endpoint)\n",
    "    2. [Unsubscribe to the listing (optional)](#B.-Unsubscribe-to-the-listing-(optional))\n",
    "    \n",
    "\n",
    "## Usage instructions\n",
    "You can run this notebook one cell at a time (By using Shift+Enter for running a cell)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Subscribe to the finetune algorithm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To subscribe to the model algorithm:\n",
    "1. Open the algorithm listing page [Cohere Command Finetuning](https://aws.amazon.com/marketplace/pp/prodview-3hsemkam64xmi)\n",
    "2. On the AWS Marketplace listing, click on the **Continue to Subscribe** button.\n",
    "3. On the **Subscribe to this software** page, review and click on **\"Accept Offer\"** if you and your organization agrees with EULA, pricing, and support terms. \n",
    "4. Once you click on **Continue to configuration** button and then choose a **region**, you will see a **Product Arn** displayed. This is the algorithm ARN that you need to specify while creating a finetune or deploying the finetuned model as an endpoint using boto3. Copy the ARN corresponding to your region and specify the same in the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!pip install -U cohere-aws\n",
    "\n",
    "from cohere_aws import Client\n",
    "import boto3\n",
    "import sagemaker as sage\n",
    "from sagemaker.s3 import S3Uploader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The algorithm is available in the list of AWS regions specified below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "region = boto3.Session().region_name\n",
    "\n",
    "cohere_package = \"cohere-command-ft-v2-1-03b5a2b33e113acdbde5f8a142f85db7\"\n",
    "\n",
    "# Mapping for algorithms\n",
    "algorithm_map = {\n",
    "    \"us-east-1\": f\"arn:aws:sagemaker:us-east-1:865070037744:algorithm/{cohere_package}\",\n",
    "    \"us-east-2\": f\"arn:aws:sagemaker:us-east-2:057799348421:algorithm/{cohere_package}\",\n",
    "    \"us-west-2\": f\"arn:aws:sagemaker:us-west-2:594846645681:algorithm/{cohere_package}\",\n",
    "    \"eu-central-1\": f\"arn:aws:sagemaker:eu-central-1:446921602837:algorithm/{cohere_package}\",\n",
    "    \"ap-southeast-1\": f\"arn:aws:sagemaker:ap-southeast-1:192199979996:algorithm/{cohere_package}\",\n",
    "    \"ap-southeast-2\": f\"arn:aws:sagemaker:ap-southeast-2:666831318237:algorithm/{cohere_package}\",\n",
    "    \"ap-northeast-1\": f\"arn:aws:sagemaker:ap-northeast-1:977537786026:algorithm/{cohere_package}\",\n",
    "    \"ap-south-1\": f\"arn:aws:sagemaker:ap-south-1:077584701553:algorithm/{cohere_package}\",\n",
    "}\n",
    "if region not in algorithm_map.keys():\n",
    "    raise Exception(f\"Current boto3 session region {region} is not supported.\")\n",
    "\n",
    "arn = algorithm_map[region]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Upload training and evaluation datasets\n",
    "\n",
    "Select a path on S3 to store the training and evaluation datasets and update the **s3_data_dir** below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_data_dir = \"s3://...\"  # Do not add a trailing slash otherwise the upload will not work"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Upload sample training data to S3:\n",
    "\n",
    "### Note:\n",
    "\n",
    "You'll need your data in a .csv or .jsonl file that contains prompt-completion pairs as your examples.\n",
    "\n",
    "\n",
    "### Example:\n",
    "\n",
    "JSONL:  `{\"prompt\": \"This is the first prompt\", \"completion\": \"This is the first completion\"}`\n",
    "\n",
    "CSV:  `\"This is the first prompt\" , \"This is the first completion\"`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sess = sage.Session()\n",
    "train_dataset = S3Uploader.upload(\"../examples/sample_sst5_finetuning_data.jsonl\", s3_data_dir, sagemaker_session=sess)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note:** Repeat the for the same for the evaluation dataset if you have one. If absent, we will auto-split the training dataset into training and evaluation datasets with the ratio of 90:10.\n",
    "\n",
    "Remember the dataset must contain at least 32 examples. If an evaluation dataset is provided, both training and evaluation datasets must contain at least 16 examples. The above split ratio is overwritten if the evaluation split is lesser than 16 examples. So for a dataset of size 50 the evaluation is 16 examples and the remaining 34 examples are used for training.\n",
    "\n",
    "We recommend using a dataset than contains at least 100 examples but a larger dataset is likely to yield high quality finetunes. Be aware that a larger dataset would mean that the time to finetune would also be larger."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### B. Finetune model on uploaded data\n",
    "\n",
    "Specify a directory on S3 where finetuned models should be stored. Make sure you do not reuse the same directory across multiple runs. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO update this with a custom S3 path\n",
    "# DO NOT re-use the same s3 directory for multiple finetunes\n",
    "# DO NOT add a trailing slash at the end\n",
    "\n",
    "s3_models_dir = \"s3://...\"  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create Cohere client:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "co = Client(region_name=region)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Optional: Define hyperparameters\n",
    "\n",
    "- `train_epochs`: This is the maximum number of training epochs to run for. Defaults to **1**.\n",
    "- `strategy`: Use either **tfew** or **vanilla**. Defaults to **tfew**, a parameter efficient finetuning approach. **vanilla** implies the weight updates will be applied to the last quarter of the layers for Command.\n",
    "- `learning_rate`: The initial learning rate to be used during training. Default differs based on ARN and strategy and is listed below.\n",
    "- `train_batch_size`: The batch size used during training. Defaults to **16** for Command.\n",
    "- `early_stopping_patience`: Stop training if the loss metric does not improve beyond 'early_stopping_threshold' for this many times of evaluation. Defaults to **6.**\n",
    "- `early_stopping_threshold`: How much the loss must improve to prevent early stopping. Defaults to **0.01**.\n",
    "\n",
    "\n",
    "| Model | Strategy | Learning Rate |\n",
    "| --- | --- | --- |\n",
    "| Command | vanilla | 1E-05 |\n",
    "| Command | tfew | 0.01 |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example of how to pass hyperparameters to the fine-tuning job\n",
    "train_parameters = {\n",
    "    \"train_epochs\": 1,\n",
    "    \"strategy\": \"tfew\",\n",
    "    \"early_stopping_patience\": 5,\n",
    "    \"early_stopping_threshold\": 0.001,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create fine-tuning jobs for the uploaded datasets. Add a field for `eval_data` if you have pre-split your dataset and uploaded both training and evaluation datasets to S3. Remember to use p4de for Command Finetuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This will take approximately 30 minutes with the example dataset\n",
    "finetune_name = \"sample-finetune\"\n",
    "co.create_finetune(arn=arn,\n",
    "    name=finetune_name,\n",
    "    train_data=train_dataset,\n",
    "    s3_models_dir=s3_models_dir,\n",
    "    instance_type=\"ml.p4de.24xlarge\",\n",
    "    training_parameters=train_parameters,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The finetuned weights for the above will be store in a tar file `{s3_models_dir}/sample-finetune.tar.gz` where the file name is the same as the name used during the creation of the finetune."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Create an endpoint for inference with the custom model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Create an endpoint\n",
    "\n",
    "The Cohere AWS SDK provides a built-in method for creating an endpoint for inference. This will automatically deploy the model you finetuned earlier.\n",
    "\n",
    "> **Note**: This is equivalent to creating and deploying a `ModelPackage` in SageMaker's SDK.\n",
    "\n",
    "You can serve multiple t-few finetunes if you store all of them in the same S3 directory and pass the directory as `s3_model_dir`. Please note that you should use dedicated directories for vanilla finetunes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "co.create_endpoint(arn=arn,\n",
    "    endpoint_name=\"command-finetune-test\",\n",
    "    s3_models_dir=s3_models_dir,\n",
    "    recreate=True,\n",
    "    instance_type=\"ml.p4d.24xlarge\",\n",
    ")\n",
    "\n",
    "# If the endpoint is already created, you just need to connect to it\n",
    "# co.connect_to_endpoint(endpoint_name=\"command-finetune-test\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### B. Perform real-time inference\n",
    "\n",
    "Now, you can access all models deployed on the endpoint for inference:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When serving t-few finetunes, you must additionally pass the `name` in the `model` parameter as specified during the creation of the finetune."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"Classify the following text as either very negative, negative, neutral, positive or very positive: mr. deeds is , as comedy goes , very silly -- and in the best way.\"\n",
    "\n",
    "# vanilla\n",
    "# result = co.generate(prompt=prompt, max_tokens=50)\n",
    "\n",
    "# tfew\n",
    "result = co.generate(model=finetune_name, prompt=prompt, max_tokens=50)\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Clean-up"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Delete the endpoint\n",
    "\n",
    "After you've successfully performed inference, you can delete the deployed endpoint to avoid being charged continuously. This can also be done via the Cohere AWS SDK:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "co.delete_endpoint()\n",
    "co.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Stacking multiple T-Few finetunes together\n",
    "\n",
    "When creating finetunes with the strategy as `tfew`, the resultant weights will be in the order of magnitude of 1-10 MB. This unlocks the ability to keep multiple T-Few finetunes in GPU VRAM by stacking them one on top of the other. This, combined with some interesting framework optimizations, allows us to perform inference for multiple T-Few finetunes concurrently. To read more about how we do stacked serving of the T-Few finetunes refer to Cohere's [T-few Finetuning blog post](https://txt.cohere.com/tfew-finetuning/).\n",
    "\n",
    "It is important to use unique names when creating `tfew` finetunes during the co.create_finetune call, so they can be stacked together. Some important notes for stacking T-Few finetunes are:\n",
    "* The T-Few finetuned weights must be created using the same version of the algorithm.\n",
    "* You must select the T-Few finetunes to be stacked together and copy the corresponding tar files to a S3 directory of your choosing. This S3 directory should then be passed to co.create_endpoint()\n",
    "* When using the SDK to create a stacked T-Few model endpoint, your collection of T-Few finetunes will be extracted and re-combined to a single tar file. It will create a `models.tar.gz` file in the same S3 directory, which is then consumed to create the Model Endpoint on Sagemaker.\n",
    "\n",
    "Lets create a second T-Few finetune so we can stack the two together and see stacked serving in action."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a second finetune\n",
    "\n",
    "Ideally use a new training dataset, but remember to use strategy of tfew, and to select a different S3 directory. Note that we also gave the finetune a different name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "finetune2_name = \"sample-finetune-v2\"\n",
    "co.create_finetune(arn=arn,\n",
    "    name=finetune2_name,\n",
    "    train_data=train_dataset,\n",
    "    s3_models_dir=s3_models_dir,\n",
    "    instance_type=\"ml.p4de.24xlarge\",\n",
    "    training_parameters=train_parameters,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once the above is complete, you can then copy both the tar files to a dedicated S3 directory, set the value for the `s3_stacked_dir` and use that to create and endpoint like before. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_stacked_dir = \"s3:/...\"\n",
    "\n",
    "co.create_endpoint(arn=arn,\n",
    "    endpoint_name=\"command-stacked-test\",\n",
    "    s3_models_dir=s3_stacked_dir,\n",
    "    recreate=True,\n",
    "    instance_type=\"ml.p4d.24xlarge\",\n",
    ")\n",
    "\n",
    "# If the endpoint is already created, you just need to connect to it\n",
    "# co.connect_to_endpoint(endpoint_name=\"command-stacked-test\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To send inference requests, use the above endpoint but select the corresponding finetune by specifying the finetune name in the `model` field. You will see it is capable of running inference for both of the created t-few finetunes. You can use this strategy to serve an arbitary number of finetunes concurrently on the same hardware."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = co.generate(model=finetune_name, prompt=prompt, max_tokens=50)\n",
    "print(result)\n",
    "\n",
    "result = co.generate(model=finetune2_name, prompt=prompt, max_tokens=50)\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want to update an existing model endpoint, such as to add or remove tfew finetunes from the stack, then you can use the [update_endpoint](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker/client/update_endpoint.html#) functionality of Sagemaker, and use rolling update as the policy to ensure there is no downtime."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Unsubscribe to the listing (optional)\n",
    "\n",
    "If you would like to unsubscribe to the model package, follow these steps. Before you cancel the subscription, ensure that you do not have any [deployable models](https://console.aws.amazon.com/sagemaker/home#/models) created from the model package or using the algorithm. Note - You can find this information by looking at the container name associated with the model. \n",
    "\n",
    "**Steps to unsubscribe to product from AWS Marketplace**:\n",
    "1. Navigate to __Machine Learning__ tab on [__Your Software subscriptions page__](https://aws.amazon.com/marketplace/ai/library?productType=ml&ref_=mlmp_gitdemo_indust)\n",
    "2. Locate the listing that you want to cancel the subscription for, and then choose __Cancel Subscription__  to cancel the subscription.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.2"
  },
  "vscode": {
   "interpreter": {
    "hash": "3b57b3736fb00bc0deb03789040183ddbda4c9eb8e8f6bef7ea4333bc64826af"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
